{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "284d3de2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !sudo apt install libopenmpi-dev -y\n",
    "# !pip3 install mpi4py --user\n",
    "# !pip3 install deepspeed==0.12.3 --user"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e4bc859d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !pip3 install accelerate transformers -U --user"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "992d1de9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "absl-py==2.0.0\r\n",
      "accelerate==0.25.0\r\n",
      "aiofiles==23.2.1\r\n",
      "aiohttp==3.8.5\r\n",
      "aiohttp-cors==0.7.0\r\n",
      "aiorwlock==1.3.0\r\n",
      "aiosignal==1.3.1\r\n",
      "altair==5.1.2\r\n",
      "anyio==3.7.1\r\n",
      "appdirs==1.4.4\r\n",
      "argon2-cffi==23.1.0\r\n",
      "argon2-cffi-bindings==21.2.0\r\n",
      "asttokens==2.2.1\r\n",
      "async-timeout==4.0.3\r\n",
      "attributedict==0.3.0\r\n",
      "attrs==23.1.0\r\n",
      "autoawq @ https://github.com/casper-hansen/AutoAWQ/releases/download/v0.1.6/autoawq-0.1.6+cu118-cp310-cp310-linux_x86_64.whl\r\n",
      "azure-core==1.29.5\r\n",
      "azure-identity==1.15.0\r\n",
      "azure-storage-blob==12.18.3\r\n",
      "azure-storage-file-datalake==12.13.2\r\n",
      "backcall==0.2.0\r\n",
      "bcrypt==4.0.1\r\n",
      "beautifulsoup4==4.12.2\r\n",
      "bitsandbytes==0.41.0\r\n",
      "bleach==6.0.0\r\n",
      "blessed==1.20.0\r\n",
      "blessings==1.7\r\n",
      "boto3==1.28.78\r\n",
      "botocore==1.31.78\r\n",
      "Brotli==1.1.0\r\n",
      "cachetools==5.3.2\r\n",
      "causal-conv1d==1.0.0\r\n",
      "certifi==2022.12.7\r\n",
      "cffi==1.15.1\r\n",
      "chardet==5.2.0\r\n",
      "charset-normalizer==2.1.1\r\n",
      "circuitbreaker==1.4.0\r\n",
      "click==8.1.7\r\n",
      "cmake==3.27.7\r\n",
      "codecov==2.1.13\r\n",
      "colorama==0.4.6\r\n",
      "coloredlogs==15.0.1\r\n",
      "colorful==0.5.5\r\n",
      "colour-runner==0.1.1\r\n",
      "comm==0.1.4\r\n",
      "contourpy==1.2.0\r\n",
      "coverage==7.3.2\r\n",
      "cryptography==41.0.5\r\n",
      "cycler==0.12.1\r\n",
      "DataProperty==1.0.1\r\n",
      "datasets==2.14.6\r\n",
      "debugpy==1.6.7.post1\r\n",
      "decorator==5.1.1\r\n",
      "deepdiff==6.6.1\r\n",
      "deepspeed==0.12.3\r\n",
      "defusedxml==0.7.1\r\n",
      "dill==0.3.7\r\n",
      "distlib==0.3.7\r\n",
      "distro==1.7.0\r\n",
      "docker-pycreds==0.4.0\r\n",
      "einops==0.6.1\r\n",
      "einops-exts==0.0.4\r\n",
      "evaluate==0.4.1\r\n",
      "exceptiongroup==1.1.3\r\n",
      "executing==1.2.0\r\n",
      "fastapi==0.104.1\r\n",
      "fastjsonschema==2.18.0\r\n",
      "ffmpy==0.3.1\r\n",
      "filelock==3.13.1\r\n",
      "flash-attn==2.3.0\r\n",
      "fonttools==4.44.0\r\n",
      "frozenlist==1.4.0\r\n",
      "fsspec==2023.10.0\r\n",
      "gitdb==4.0.11\r\n",
      "GitPython==3.1.40\r\n",
      "google-api-core==2.12.0\r\n",
      "google-auth==2.23.4\r\n",
      "google-cloud-core==2.3.3\r\n",
      "google-cloud-storage==2.10.0\r\n",
      "google-crc32c==1.5.0\r\n",
      "google-resumable-media==2.6.0\r\n",
      "googleapis-common-protos==1.61.0\r\n",
      "gpustat==1.1.1\r\n",
      "gradio==3.35.2\r\n",
      "gradio_client==0.2.9\r\n",
      "grpcio==1.59.2\r\n",
      "h11==0.14.0\r\n",
      "hjson==3.1.0\r\n",
      "httpcore==0.17.3\r\n",
      "httptools==0.6.1\r\n",
      "httpx==0.24.0\r\n",
      "huggingface-hub==0.17.3\r\n",
      "humanfriendly==10.0\r\n",
      "idna==3.4\r\n",
      "inspecta==0.1.3\r\n",
      "ipykernel==6.25.1\r\n",
      "ipython==8.14.0\r\n",
      "ipython-genutils==0.2.0\r\n",
      "ipywidgets==8.1.0\r\n",
      "isodate==0.6.1\r\n",
      "jedi==0.19.0\r\n",
      "Jinja2==3.1.2\r\n",
      "jmespath==1.0.1\r\n",
      "joblib==1.3.2\r\n",
      "jsonlines==4.0.0\r\n",
      "jsonschema==4.19.0\r\n",
      "jsonschema-specifications==2023.7.1\r\n",
      "jupyter==1.0.0\r\n",
      "jupyter-console==6.6.3\r\n",
      "jupyter-server==1.18.0\r\n",
      "jupyter-server-proxy==3.2.1\r\n",
      "jupyter_client==8.3.0\r\n",
      "jupyter_core==5.3.1\r\n",
      "jupyterlab-pygments==0.2.2\r\n",
      "jupyterlab-widgets==3.0.8\r\n",
      "kiwisolver==1.4.5\r\n",
      "linkify-it-py==2.0.2\r\n",
      "lit==17.0.4\r\n",
      "# Editable install with no version control (llava==1.1.3)\r\n",
      "-e /home/ubuntu/LLaVA\r\n",
      "lm-eval==0.3.0\r\n",
      "mamba-ssm @ file:///home/ubuntu/mamba\r\n",
      "markdown-it-py==2.2.0\r\n",
      "markdown2==2.4.10\r\n",
      "MarkupSafe==2.1.3\r\n",
      "matplotlib==3.8.1\r\n",
      "matplotlib-inline==0.1.6\r\n",
      "mbstrdecoder==1.1.3\r\n",
      "mdit-py-plugins==0.3.3\r\n",
      "mdurl==0.1.2\r\n",
      "mistune==3.0.1\r\n",
      "mosaicml-streaming==0.6.1\r\n",
      "mp==0.5.0\r\n",
      "mpi4py==3.1.5\r\n",
      "mpmath==1.3.0\r\n",
      "msal==1.25.0\r\n",
      "msal-extensions==1.0.0\r\n",
      "msgpack==1.0.7\r\n",
      "msgspec==0.18.4\r\n",
      "multidict==6.0.4\r\n",
      "multiprocess==0.70.15\r\n",
      "nbclient==0.8.0\r\n",
      "nbconvert==7.7.4\r\n",
      "nbformat==5.9.2\r\n",
      "nest-asyncio==1.5.7\r\n",
      "networkx==3.0\r\n",
      "ninja==1.11.1.1\r\n",
      "nltk==3.8.1\r\n",
      "notebook==6.4.12\r\n",
      "numexpr==2.8.7\r\n",
      "numpy==1.24.1\r\n",
      "nvidia-cublas-cu11==11.10.3.66\r\n",
      "nvidia-cuda-cupti-cu11==11.7.101\r\n",
      "nvidia-cuda-nvrtc-cu11==11.7.99\r\n",
      "nvidia-cuda-runtime-cu11==11.7.99\r\n",
      "nvidia-cudnn-cu11==8.5.0.96\r\n",
      "nvidia-cufft-cu11==10.9.0.58\r\n",
      "nvidia-curand-cu11==10.2.10.91\r\n",
      "nvidia-cusolver-cu11==11.4.0.1\r\n",
      "nvidia-cusparse-cu11==11.7.4.91\r\n",
      "nvidia-ml-py==12.535.133\r\n",
      "nvidia-nccl-cu11==2.14.3\r\n",
      "nvidia-nvtx-cu11==11.7.91\r\n",
      "oci==2.115.0\r\n",
      "openai==0.28.0\r\n",
      "opencensus==0.11.3\r\n",
      "opencensus-context==0.1.3\r\n",
      "ordered-set==4.1.0\r\n",
      "orjson==3.9.10\r\n",
      "packaging==23.1\r\n",
      "pandas==2.1.2\r\n",
      "pandocfilters==1.5.0\r\n",
      "paramiko==3.3.1\r\n",
      "parso==0.8.3\r\n",
      "pathvalidate==3.2.0\r\n",
      "peft==0.4.0\r\n",
      "pexpect==4.8.0\r\n",
      "pickleshare==0.7.5\r\n",
      "Pillow==9.3.0\r\n",
      "platformdirs==3.10.0\r\n",
      "pluggy==1.3.0\r\n",
      "portalocker==2.8.2\r\n",
      "prometheus-client==0.17.1\r\n",
      "prompt-toolkit==3.0.39\r\n",
      "protobuf==4.25.0\r\n",
      "psutil==5.9.5\r\n",
      "ptyprocess==0.7.0\r\n",
      "pure-eval==0.2.2\r\n",
      "py-cpuinfo==9.0.0\r\n",
      "py-spy==0.3.14\r\n",
      "pyarrow==14.0.0\r\n",
      "pyasn1==0.5.0\r\n",
      "pyasn1-modules==0.3.0\r\n",
      "pybind11==2.11.1\r\n",
      "pycountry==22.3.5\r\n",
      "pycparser==2.21\r\n",
      "pydantic==1.10.13\r\n",
      "pydub==0.25.1\r\n",
      "Pygments==2.16.1\r\n",
      "PyJWT==2.8.0\r\n",
      "PyNaCl==1.5.0\r\n",
      "pynvml==11.5.0\r\n",
      "pyOpenSSL==23.3.0\r\n",
      "pyparsing==3.1.1\r\n",
      "pyproject-api==1.6.1\r\n",
      "pytablewriter==1.2.0\r\n",
      "python-dateutil==2.8.2\r\n",
      "python-dotenv==1.0.0\r\n",
      "python-multipart==0.0.6\r\n",
      "python-snappy==0.6.1\r\n",
      "pytz==2023.3.post1\r\n",
      "PyYAML==6.0.1\r\n",
      "pyzmq==25.1.1\r\n",
      "qtconsole==5.4.3\r\n",
      "QtPy==2.3.1\r\n",
      "ray==2.8.0\r\n",
      "referencing==0.30.2\r\n",
      "regex==2023.10.3\r\n",
      "requests==2.28.1\r\n",
      "responses==0.18.0\r\n",
      "rich==13.6.0\r\n",
      "rootpath==0.1.1\r\n",
      "rouge-score==0.1.2\r\n",
      "rpds-py==0.9.2\r\n",
      "rsa==4.9\r\n",
      "s3transfer==0.7.0\r\n",
      "sacrebleu==1.5.0\r\n",
      "safetensors==0.4.0\r\n",
      "scikit-learn==1.2.2\r\n",
      "scipy==1.11.3\r\n",
      "semantic-version==2.10.0\r\n",
      "Send2Trash==1.8.2\r\n",
      "sentencepiece==0.1.99\r\n",
      "sentry-sdk==1.36.0\r\n",
      "setproctitle==1.3.3\r\n",
      "shortuuid==1.0.11\r\n",
      "simpervisor==1.0.0\r\n",
      "six==1.16.0\r\n",
      "smart-open==6.4.0\r\n",
      "smmap==5.0.1\r\n",
      "sniffio==1.3.0\r\n",
      "soupsieve==2.4.1\r\n",
      "sqlitedict==2.1.0\r\n",
      "ssh-import-id==5.11\r\n",
      "stack-data==0.6.2\r\n",
      "starlette==0.27.0\r\n",
      "svgwrite==1.4.3\r\n",
      "sympy==1.12\r\n",
      "tabledata==1.3.3\r\n",
      "tabulate==0.9.0\r\n",
      "tcolorpy==0.1.4\r\n",
      "tensorboardX==2.6.2.2\r\n",
      "termcolor==2.3.0\r\n",
      "terminado==0.17.1\r\n",
      "texttable==1.7.0\r\n",
      "threadpoolctl==3.2.0\r\n",
      "timm==0.6.13\r\n",
      "tinycss2==1.2.1\r\n",
      "tokenizers==0.14.1\r\n",
      "toml==0.10.2\r\n",
      "tomli==2.0.1\r\n",
      "toolz==0.12.0\r\n",
      "torch==2.0.1+cu118\r\n",
      "torchaudio==2.0.2+cu118\r\n",
      "torchvision==0.15.2+cu118\r\n",
      "tornado==6.3.3\r\n",
      "tox==4.11.3\r\n",
      "tqdm==4.66.1\r\n",
      "tqdm-multiprocess==0.0.11\r\n",
      "traitlets==5.9.0\r\n",
      "transformers==4.35.2\r\n",
      "triton==2.0.0\r\n",
      "typepy==1.3.2\r\n",
      "typing_extensions==4.8.0\r\n",
      "tzdata==2023.3\r\n",
      "uc-micro-py==1.0.2\r\n",
      "unzip==1.0.0\r\n",
      "urllib3==1.26.13\r\n",
      "uvicorn==0.24.0.post1\r\n",
      "uvloop==0.19.0\r\n",
      "virtualenv==20.21.0\r\n",
      "wandb==0.16.0\r\n",
      "watchfiles==0.21.0\r\n",
      "wavedrom==2.0.3.post3\r\n",
      "wcwidth==0.2.6\r\n",
      "webencodings==0.5.1\r\n",
      "websocket-client==1.6.1\r\n",
      "websockets==12.0\r\n",
      "widgetsnbextension==4.0.8\r\n",
      "xxhash==3.4.1\r\n",
      "yarl==1.9.2\r\n",
      "zstandard==0.22.0\r\n",
      "zstd==1.5.5.1\r\n"
     ]
    }
   ],
   "source": [
    "!pip3 freeze"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ecf642c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tue Dec  5 08:51:21 2023       \r\n",
      "+-----------------------------------------------------------------------------+\r\n",
      "| NVIDIA-SMI 525.85.12    Driver Version: 525.85.12    CUDA Version: 12.0     |\r\n",
      "|-------------------------------+----------------------+----------------------+\r\n",
      "| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |\r\n",
      "| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |\r\n",
      "|                               |                      |               MIG M. |\r\n",
      "|===============================+======================+======================|\r\n",
      "|   0  NVIDIA A100 80G...  On   | 00000001:00:00.0 Off |                    0 |\r\n",
      "| N/A   38C    P0    63W / 300W |  22394MiB / 81920MiB |      0%      Default |\r\n",
      "|                               |                      |             Disabled |\r\n",
      "+-------------------------------+----------------------+----------------------+\r\n",
      "                                                                               \r\n",
      "+-----------------------------------------------------------------------------+\r\n",
      "| Processes:                                                                  |\r\n",
      "|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |\r\n",
      "|        ID   ID                                                   Usage      |\r\n",
      "|=============================================================================|\r\n",
      "+-----------------------------------------------------------------------------+\r\n"
     ]
    }
   ],
   "source": [
    "!nvidia-smi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b2374da5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"WANDB_DISABLED\"] = \"true\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "9b310aa0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (c) 2023, Albert Gu, Tri Dao.\n",
    "\n",
    "import math\n",
    "from functools import partial\n",
    "\n",
    "from collections import namedtuple\n",
    "\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "from torch.nn import CrossEntropyLoss\n",
    "from transformers import PretrainedConfig\n",
    "\n",
    "from mamba_ssm.modules.mamba_simple import Mamba, Block\n",
    "from mamba_ssm.utils.generation import GenerationMixin\n",
    "from mamba_ssm.utils.hf import load_config_hf, load_state_dict_hf\n",
    "\n",
    "try:\n",
    "    from mamba_ssm.ops.triton.layernorm import RMSNorm, layer_norm_fn, rms_norm_fn\n",
    "except ImportError:\n",
    "    RMSNorm, layer_norm_fn, rms_norm_fn = None, None, None\n",
    "\n",
    "\n",
    "def create_block(\n",
    "    d_model,\n",
    "    ssm_cfg=None,\n",
    "    norm_epsilon=1e-5,\n",
    "    rms_norm=False,\n",
    "    residual_in_fp32=False,\n",
    "    fused_add_norm=False,\n",
    "    layer_idx=None,\n",
    "    device=None,\n",
    "    dtype=None,\n",
    "):\n",
    "    if ssm_cfg is None:\n",
    "        ssm_cfg = {}\n",
    "    factory_kwargs = {\"device\": device, \"dtype\": dtype}\n",
    "    mixer_cls = partial(Mamba, layer_idx=layer_idx, **ssm_cfg, **factory_kwargs)\n",
    "    norm_cls = partial(\n",
    "        nn.LayerNorm if not rms_norm else RMSNorm, eps=norm_epsilon, **factory_kwargs\n",
    "    )\n",
    "    block = Block(\n",
    "        d_model,\n",
    "        mixer_cls,\n",
    "        norm_cls=norm_cls,\n",
    "        fused_add_norm=fused_add_norm,\n",
    "        residual_in_fp32=residual_in_fp32,\n",
    "    )\n",
    "    block.layer_idx = layer_idx\n",
    "    return block\n",
    "\n",
    "\n",
    "# https://github.com/huggingface/transformers/blob/c28d04e9e252a1a099944e325685f14d242ecdcd/src/transformers/models/gpt2/modeling_gpt2.py#L454\n",
    "def _init_weights(\n",
    "    module,\n",
    "    n_layer,\n",
    "    initializer_range=0.02,  # Now only used for embedding layer.\n",
    "    rescale_prenorm_residual=True,\n",
    "    n_residuals_per_layer=1,  # Change to 2 if we have MLP\n",
    "):\n",
    "    if isinstance(module, nn.Linear):\n",
    "        if module.bias is not None:\n",
    "            if not getattr(module.bias, \"_no_reinit\", False):\n",
    "                nn.init.zeros_(module.bias)\n",
    "    elif isinstance(module, nn.Embedding):\n",
    "        nn.init.normal_(module.weight, std=initializer_range)\n",
    "\n",
    "    if rescale_prenorm_residual:\n",
    "        # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:\n",
    "        #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale\n",
    "        #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.\n",
    "        #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/\n",
    "        #\n",
    "        # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py\n",
    "        for name, p in module.named_parameters():\n",
    "            if name in [\"out_proj.weight\", \"fc2.weight\"]:\n",
    "                # Special Scaled Initialization --> There are 2 Layer Norms per Transformer Block\n",
    "                # Following Pytorch init, except scale by 1/sqrt(2 * n_layer)\n",
    "                # We need to reinit p since this code could be called multiple times\n",
    "                # Having just p *= scale would repeatedly scale it down\n",
    "                nn.init.kaiming_uniform_(p, a=math.sqrt(5))\n",
    "                with torch.no_grad():\n",
    "                    p /= math.sqrt(n_residuals_per_layer * n_layer)\n",
    "\n",
    "\n",
    "class MixerModel(nn.Module):\n",
    "    def __init__(\n",
    "        self,\n",
    "        d_model: int,\n",
    "        n_layer: int,\n",
    "        vocab_size: int,\n",
    "        ssm_cfg=None,\n",
    "        norm_epsilon: float = 1e-5,\n",
    "        rms_norm: bool = False,\n",
    "        initializer_cfg=None,\n",
    "        fused_add_norm=False,\n",
    "        residual_in_fp32=False,\n",
    "        device=None,\n",
    "        dtype=None,\n",
    "    ) -> None:\n",
    "        factory_kwargs = {\"device\": device, \"dtype\": dtype}\n",
    "        super().__init__()\n",
    "        self.residual_in_fp32 = residual_in_fp32\n",
    "\n",
    "        self.embedding = nn.Embedding(vocab_size, d_model, **factory_kwargs)\n",
    "\n",
    "        # We change the order of residual and layer norm:\n",
    "        # Instead of LN -> Attn / MLP -> Add, we do:\n",
    "        # Add -> LN -> Attn / MLP / Mixer, returning both the residual branch (output of Add) and\n",
    "        # the main branch (output of MLP / Mixer). The model definition is unchanged.\n",
    "        # This is for performance reason: we can fuse add + layer_norm.\n",
    "        self.fused_add_norm = fused_add_norm\n",
    "        if self.fused_add_norm:\n",
    "            if layer_norm_fn is None or rms_norm_fn is None:\n",
    "                raise ImportError(\"Failed to import Triton LayerNorm / RMSNorm kernels\")\n",
    "\n",
    "        self.layers = nn.ModuleList(\n",
    "            [\n",
    "                create_block(\n",
    "                    d_model,\n",
    "                    ssm_cfg=ssm_cfg,\n",
    "                    norm_epsilon=norm_epsilon,\n",
    "                    rms_norm=rms_norm,\n",
    "                    residual_in_fp32=residual_in_fp32,\n",
    "                    fused_add_norm=fused_add_norm,\n",
    "                    layer_idx=i,\n",
    "                    **factory_kwargs,\n",
    "                )\n",
    "                for i in range(n_layer)\n",
    "            ]\n",
    "        )\n",
    "\n",
    "        self.norm_f = (nn.LayerNorm if not rms_norm else RMSNorm)(\n",
    "            d_model, eps=norm_epsilon, **factory_kwargs\n",
    "        )\n",
    "\n",
    "        self.apply(\n",
    "            partial(\n",
    "                _init_weights,\n",
    "                n_layer=n_layer,\n",
    "                **(initializer_cfg if initializer_cfg is not None else {}),\n",
    "            )\n",
    "        )\n",
    "\n",
    "    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):\n",
    "        return {\n",
    "            i: layer.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)\n",
    "            for i, layer in enumerate(self.layers)\n",
    "        }\n",
    "\n",
    "    def forward(self, input_ids, inference_params=None):\n",
    "        hidden_states = self.embedding(input_ids)\n",
    "        residual = None\n",
    "        for layer in self.layers:\n",
    "            hidden_states, residual = layer(\n",
    "                hidden_states, residual, inference_params=inference_params\n",
    "            )\n",
    "        if not self.fused_add_norm:\n",
    "            residual = (hidden_states + residual) if residual is not None else hidden_states\n",
    "            hidden_states = self.norm_f(residual.to(dtype=self.norm_f.weight.dtype))\n",
    "        else:\n",
    "            # Set prenorm=False here since we don't need the residual\n",
    "            fused_add_norm_fn = rms_norm_fn if isinstance(self.norm_f, RMSNorm) else layer_norm_fn\n",
    "            hidden_states = fused_add_norm_fn(\n",
    "                hidden_states,\n",
    "                self.norm_f.weight,\n",
    "                self.norm_f.bias,\n",
    "                eps=self.norm_f.eps,\n",
    "                residual=residual,\n",
    "                prenorm=False,\n",
    "                residual_in_fp32=self.residual_in_fp32,\n",
    "            )\n",
    "        return hidden_states\n",
    "\n",
    "\n",
    "class MambaLMHeadModel(nn.Module, GenerationMixin):\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        d_model: int,\n",
    "        n_layer: int,\n",
    "        vocab_size: int,\n",
    "        initializer_cfg=None,\n",
    "        pad_vocab_size_multiple: int = 1,\n",
    "        device=None,\n",
    "        dtype=None,\n",
    "        **backbone_kwargs,\n",
    "    ) -> None:\n",
    "        factory_kwargs = {\"device\": device, \"dtype\": dtype}\n",
    "        super().__init__()\n",
    "        if vocab_size % pad_vocab_size_multiple != 0:\n",
    "            vocab_size += pad_vocab_size_multiple - (vocab_size % pad_vocab_size_multiple)\n",
    "        self.backbone = MixerModel(\n",
    "            d_model=d_model,\n",
    "            n_layer=n_layer,\n",
    "            vocab_size=vocab_size,\n",
    "            initializer_cfg=initializer_cfg,\n",
    "            **backbone_kwargs,\n",
    "            **factory_kwargs,\n",
    "        )\n",
    "        self.lm_head = nn.Linear(d_model, vocab_size, bias=False, **factory_kwargs)\n",
    "\n",
    "        # Initialize weights and apply final processing\n",
    "        self.apply(\n",
    "            partial(\n",
    "                _init_weights,\n",
    "                n_layer=n_layer,\n",
    "                **(initializer_cfg if initializer_cfg is not None else {}),\n",
    "            )\n",
    "        )\n",
    "        self.tie_weights()\n",
    "        self.config = PretrainedConfig(\n",
    "            d_model = d_model,\n",
    "            n_layer = n_layer,\n",
    "            vocab_size = vocab_size,\n",
    "            hidden_size = d_model,\n",
    "        )\n",
    "\n",
    "    def tie_weights(self):\n",
    "        self.lm_head.weight = self.backbone.embedding.weight\n",
    "\n",
    "    def allocate_inference_cache(self, batch_size, max_seqlen, dtype=None, **kwargs):\n",
    "        return self.backbone.allocate_inference_cache(batch_size, max_seqlen, dtype=dtype, **kwargs)\n",
    "\n",
    "    def forward(self, input_ids, position_ids=None, inference_params=None, num_last_tokens=0, labels = None):\n",
    "        \"\"\"\n",
    "        \"position_ids\" is just to be compatible with Transformer generation. We don't use it.\n",
    "        num_last_tokens: if > 0, only return the logits for the last n tokens\n",
    "        \"\"\"\n",
    "        hidden_states = self.backbone(input_ids, inference_params=inference_params)\n",
    "        if num_last_tokens > 0:\n",
    "            hidden_states = hidden_states[:, -num_last_tokens:]\n",
    "        lm_logits = self.lm_head(hidden_states)\n",
    "        \n",
    "        loss = None\n",
    "        if labels is not None:\n",
    "            logits = lm_logits\n",
    "            # Shift so that tokens < n predict n\n",
    "            shift_logits = logits[..., :-1, :].contiguous()\n",
    "            shift_labels = labels[..., 1:].contiguous()\n",
    "            # Flatten the tokens\n",
    "            loss_fct = CrossEntropyLoss()\n",
    "            shift_logits = shift_logits.view(-1, self.config.vocab_size)\n",
    "            shift_labels = shift_labels.view(-1)\n",
    "            # Enable model parallelism\n",
    "            shift_labels = shift_labels.to(shift_logits.device)\n",
    "            loss = loss_fct(shift_logits, shift_labels)\n",
    "            print(loss, shift_logits, shift_logits.dtype, shift_labels, shift_labels.dtype)\n",
    "            return (loss,)\n",
    "            \n",
    "        else:\n",
    "            CausalLMOutput = namedtuple(\"CausalLMOutput\", [\"logits\"])\n",
    "            return CausalLMOutput(logits=lm_logits)\n",
    "\n",
    "    @classmethod\n",
    "    def from_pretrained(cls, pretrained_model_name, device=None, dtype=None, **kwargs):\n",
    "        config = load_config_hf(pretrained_model_name)\n",
    "        model = cls(**config, device=device, dtype=dtype, **kwargs)\n",
    "        model.load_state_dict(load_state_dict_hf(pretrained_model_name, device=device, dtype=dtype))\n",
    "        return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "4fe72f34",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !wget https://huggingface.co/state-spaces/mamba-130m/raw/main/config.json -O config-130m.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "eac72d2b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import json\n",
    "\n",
    "with open('config-130m.json') as fopen:\n",
    "    config = json.load(fopen)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "b85b48be",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = MambaLMHeadModel(**{**config, 'vocab_size': 32000})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "4271a418",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "32000"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.config.vocab_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "e8beea8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from streaming import LocalDataset\n",
    "import numpy as np\n",
    "from streaming.base.format.mds.encodings import Encoding, _encodings\n",
    "\n",
    "class UInt16(Encoding):\n",
    "    def encode(self, obj) -> bytes:\n",
    "        return obj.tobytes()\n",
    "\n",
    "    def decode(self, data: bytes):\n",
    "        return np.frombuffer(data, np.uint16)\n",
    "\n",
    "_encodings['uint16'] = UInt16"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "74c0f1fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# !git lfs clone https://huggingface.co/datasets/malaysia-ai/mosaic-instructions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "88f14bc3",
   "metadata": {},
   "outputs": [],
   "source": [
    "class DatasetFixed(torch.utils.data.Dataset):\n",
    "    def __init__(self, local):\n",
    "        self.dataset = LocalDataset(local=local)\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        print(idx)\n",
    "        data = self.dataset[idx]\n",
    "        data['labels'] = data['input_ids'].copy()\n",
    "\n",
    "        data.pop('token_type_ids', None)\n",
    "        for k in data.keys():\n",
    "            data[k] = data[k].astype(np.int64)\n",
    "        return data\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.dataset)\n",
    "\n",
    "train_dataset = DatasetFixed(local='mosaic-instructions')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "7ea1956b",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Using the `WANDB_DISABLED` environment variable is deprecated and will be removed in v5. Use the --report_to flag to control the integrations used for logging result (for instance --report_to none).\n"
     ]
    }
   ],
   "source": [
    "from transformers import TrainingArguments, Trainer, default_data_collator\n",
    "\n",
    "output_dir = 'test-130m'\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir,\n",
    "    per_device_train_batch_size=2,\n",
    "    gradient_accumulation_steps=1,\n",
    "    logging_steps=1,\n",
    "    save_strategy='steps',\n",
    "    save_steps=100,\n",
    "    num_train_epochs=3,\n",
    "    learning_rate=1e-4,\n",
    "    weight_decay=0,\n",
    "    warmup_steps=1000,\n",
    "    bf16=True,\n",
    "    fp16=False,\n",
    "    gradient_checkpointing=False,\n",
    "    save_total_limit=5,\n",
    "    log_level='debug',\n",
    "    max_steps=10,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "d3e45910",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "max_steps is given, it will override any value given in num_train_epochs\n",
      "Using auto half precision backend\n"
     ]
    }
   ],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=training_args,\n",
    "    train_dataset=train_dataset,\n",
    "    data_collator=default_data_collator,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "98e0c6d2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'args': TrainingArguments(\n",
       " _n_gpu=1,\n",
       " adafactor=False,\n",
       " adam_beta1=0.9,\n",
       " adam_beta2=0.999,\n",
       " adam_epsilon=1e-08,\n",
       " auto_find_batch_size=False,\n",
       " bf16=True,\n",
       " bf16_full_eval=False,\n",
       " data_seed=None,\n",
       " dataloader_drop_last=False,\n",
       " dataloader_num_workers=0,\n",
       " dataloader_pin_memory=True,\n",
       " ddp_backend=None,\n",
       " ddp_broadcast_buffers=None,\n",
       " ddp_bucket_cap_mb=None,\n",
       " ddp_find_unused_parameters=None,\n",
       " ddp_timeout=1800,\n",
       " debug=[],\n",
       " deepspeed=None,\n",
       " disable_tqdm=False,\n",
       " dispatch_batches=None,\n",
       " do_eval=False,\n",
       " do_predict=False,\n",
       " do_train=False,\n",
       " eval_accumulation_steps=None,\n",
       " eval_delay=0,\n",
       " eval_steps=None,\n",
       " evaluation_strategy=no,\n",
       " fp16=False,\n",
       " fp16_backend=auto,\n",
       " fp16_full_eval=False,\n",
       " fp16_opt_level=O1,\n",
       " fsdp=[],\n",
       " fsdp_config={'min_num_params': 0, 'xla': False, 'xla_fsdp_grad_ckpt': False},\n",
       " fsdp_min_num_params=0,\n",
       " fsdp_transformer_layer_cls_to_wrap=None,\n",
       " full_determinism=False,\n",
       " gradient_accumulation_steps=1,\n",
       " gradient_checkpointing=False,\n",
       " gradient_checkpointing_kwargs=None,\n",
       " greater_is_better=None,\n",
       " group_by_length=False,\n",
       " half_precision_backend=auto,\n",
       " hub_always_push=False,\n",
       " hub_model_id=None,\n",
       " hub_private_repo=False,\n",
       " hub_strategy=every_save,\n",
       " hub_token=<HUB_TOKEN>,\n",
       " ignore_data_skip=False,\n",
       " include_inputs_for_metrics=False,\n",
       " include_tokens_per_second=False,\n",
       " jit_mode_eval=False,\n",
       " label_names=None,\n",
       " label_smoothing_factor=0.0,\n",
       " learning_rate=0.0001,\n",
       " length_column_name=length,\n",
       " load_best_model_at_end=False,\n",
       " local_rank=0,\n",
       " log_level=debug,\n",
       " log_level_replica=warning,\n",
       " log_on_each_node=True,\n",
       " logging_dir=test-130m/runs/Dec05_08-51-41_jupyter-0,\n",
       " logging_first_step=False,\n",
       " logging_nan_inf_filter=True,\n",
       " logging_steps=1,\n",
       " logging_strategy=steps,\n",
       " lr_scheduler_type=linear,\n",
       " max_grad_norm=1.0,\n",
       " max_steps=10,\n",
       " metric_for_best_model=None,\n",
       " mp_parameters=,\n",
       " neftune_noise_alpha=None,\n",
       " no_cuda=False,\n",
       " num_train_epochs=3,\n",
       " optim=adamw_torch,\n",
       " optim_args=None,\n",
       " output_dir=test-130m,\n",
       " overwrite_output_dir=False,\n",
       " past_index=-1,\n",
       " per_device_eval_batch_size=8,\n",
       " per_device_train_batch_size=2,\n",
       " prediction_loss_only=False,\n",
       " push_to_hub=False,\n",
       " push_to_hub_model_id=None,\n",
       " push_to_hub_organization=None,\n",
       " push_to_hub_token=<PUSH_TO_HUB_TOKEN>,\n",
       " ray_scope=last,\n",
       " remove_unused_columns=True,\n",
       " report_to=['tensorboard'],\n",
       " resume_from_checkpoint=None,\n",
       " run_name=test-130m,\n",
       " save_on_each_node=False,\n",
       " save_safetensors=True,\n",
       " save_steps=100,\n",
       " save_strategy=steps,\n",
       " save_total_limit=5,\n",
       " seed=42,\n",
       " skip_memory_metrics=True,\n",
       " split_batches=False,\n",
       " tf32=None,\n",
       " torch_compile=False,\n",
       " torch_compile_backend=None,\n",
       " torch_compile_mode=None,\n",
       " torchdynamo=None,\n",
       " tpu_metrics_debug=False,\n",
       " tpu_num_cores=None,\n",
       " use_cpu=False,\n",
       " use_ipex=False,\n",
       " use_legacy_prediction_loop=False,\n",
       " use_mps_device=False,\n",
       " warmup_ratio=0.0,\n",
       " warmup_steps=1000,\n",
       " weight_decay=0,\n",
       " ),\n",
       " 'hp_name': None,\n",
       " 'deepspeed': None,\n",
       " 'is_in_train': False,\n",
       " 'accelerator': <accelerate.accelerator.Accelerator at 0x7f25ed0f6080>,\n",
       " 'is_deepspeed_enabled': False,\n",
       " 'is_fsdp_enabled': False,\n",
       " '_memory_tracker': <transformers.trainer_utils.TrainerMemoryTracker at 0x7f25ee4605e0>,\n",
       " 'model_init': None,\n",
       " 'is_model_parallel': False,\n",
       " 'fsdp': None,\n",
       " 'place_model_on_device': True,\n",
       " 'data_collator': <function transformers.data.data_collator.default_data_collator(features: List[transformers.data.data_collator.InputDataClass], return_tensors='pt') -> Dict[str, Any]>,\n",
       " 'train_dataset': <__main__.DatasetFixed at 0x7f249cddab60>,\n",
       " 'eval_dataset': None,\n",
       " 'tokenizer': None,\n",
       " 'model_wrapped': MambaLMHeadModel(\n",
       "   (backbone): MixerModel(\n",
       "     (embedding): Embedding(32000, 768)\n",
       "     (layers): ModuleList(\n",
       "       (0-23): 24 x Block(\n",
       "         (mixer): Mamba(\n",
       "           (in_proj): Linear(in_features=768, out_features=3072, bias=False)\n",
       "           (conv1d): Conv1d(1536, 1536, kernel_size=(4,), stride=(1,), padding=(3,), groups=1536)\n",
       "           (act): SiLU()\n",
       "           (x_proj): Linear(in_features=1536, out_features=80, bias=False)\n",
       "           (dt_proj): Linear(in_features=48, out_features=1536, bias=True)\n",
       "           (out_proj): Linear(in_features=1536, out_features=768, bias=False)\n",
       "         )\n",
       "         (norm): RMSNorm()\n",
       "       )\n",
       "     )\n",
       "     (norm_f): RMSNorm()\n",
       "   )\n",
       "   (lm_head): Linear(in_features=768, out_features=32000, bias=False)\n",
       " ),\n",
       " 'model': MambaLMHeadModel(\n",
       "   (backbone): MixerModel(\n",
       "     (embedding): Embedding(32000, 768)\n",
       "     (layers): ModuleList(\n",
       "       (0-23): 24 x Block(\n",
       "         (mixer): Mamba(\n",
       "           (in_proj): Linear(in_features=768, out_features=3072, bias=False)\n",
       "           (conv1d): Conv1d(1536, 1536, kernel_size=(4,), stride=(1,), padding=(3,), groups=1536)\n",
       "           (act): SiLU()\n",
       "           (x_proj): Linear(in_features=1536, out_features=80, bias=False)\n",
       "           (dt_proj): Linear(in_features=48, out_features=1536, bias=True)\n",
       "           (out_proj): Linear(in_features=1536, out_features=768, bias=False)\n",
       "         )\n",
       "         (norm): RMSNorm()\n",
       "       )\n",
       "     )\n",
       "     (norm_f): RMSNorm()\n",
       "   )\n",
       "   (lm_head): Linear(in_features=768, out_features=32000, bias=False)\n",
       " ),\n",
       " 'neftune_noise_alpha': None,\n",
       " 'compute_metrics': None,\n",
       " 'preprocess_logits_for_metrics': None,\n",
       " 'optimizer': None,\n",
       " 'lr_scheduler': None,\n",
       " 'callback_handler': <transformers.trainer_callback.CallbackHandler at 0x7f249c9e2200>,\n",
       " '_loggers_initialized': False,\n",
       " 'hub_model_id': None,\n",
       " '_signature_columns': None,\n",
       " 'use_apex': False,\n",
       " 'use_cpu_amp': False,\n",
       " 'label_smoother': None,\n",
       " 'state': TrainerState(epoch=None, global_step=0, max_steps=0, logging_steps=500, eval_steps=500, save_steps=500, num_train_epochs=0, total_flos=0, log_history=[], best_metric=None, best_model_checkpoint=None, is_local_process_zero=True, is_world_process_zero=True, is_hyper_param_search=False, trial_name=None, trial_params=None),\n",
       " 'control': TrainerControl(should_training_stop=False, should_epoch_stop=False, should_save=False, should_evaluate=False, should_log=False),\n",
       " 'current_flos': 0,\n",
       " 'hp_search_backend': None,\n",
       " 'use_tune_checkpoints': False,\n",
       " 'label_names': ['labels'],\n",
       " 'can_return_loss': False,\n",
       " '_train_batch_size': 2,\n",
       " '_created_lr_scheduler': False}"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.__dict__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "f1c6775c",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Currently training with a batch size of: 2\n",
      "***** Running training *****\n",
      "  Num examples = 385,224\n",
      "  Num Epochs = 1\n",
      "  Instantaneous batch size per device = 2\n",
      "  Total train batch size (w. parallel, distributed & accumulation) = 2\n",
      "  Gradient Accumulation steps = 1\n",
      "  Total optimization steps = 10\n",
      "  Number of trainable parameters = 115,096,320\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "327342\n",
      "345828\n",
      "318900\n",
      "69472\n",
      "tensor(10.5414, device='cuda:0', grad_fn=<NllLossBackward0>) tensor([[-0.4297,  0.5625,  0.3047,  ...,  0.1758,  0.2139,  0.1147],\n",
      "        [-0.6641,  0.0525, -0.3242,  ...,  0.1777,  0.3535,  0.3770],\n",
      "        [-0.4766, -0.1094,  0.0913,  ..., -0.7891, -0.2637, -0.4746],\n",
      "        ...,\n",
      "        [-0.0366, -0.5781, -0.9766,  ...,  0.8789,  0.9688, -0.1875],\n",
      "        [-0.3145,  0.3809,  0.3086,  ..., -0.8320, -0.3984, -0.2041],\n",
      "        [ 0.2246, -0.2061, -0.6562,  ..., -0.6250, -0.9219,  0.3984]],\n",
      "       device='cuda:0', dtype=torch.bfloat16, grad_fn=<ViewBackward0>) torch.bfloat16 tensor([ 6845,  5341,  3474,  ...,    11, 15119,  4318], device='cuda:0') torch.int64\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "\n",
       "    <div>\n",
       "      \n",
       "      <progress value='10' max='10' style='width:300px; height:20px; vertical-align: middle;'></progress>\n",
       "      [10/10 00:01, Epoch 0/1]\n",
       "    </div>\n",
       "    <table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       " <tr style=\"text-align: left;\">\n",
       "      <th>Step</th>\n",
       "      <th>Training Loss</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <td>1</td>\n",
       "      <td>10.541400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>2</td>\n",
       "      <td>10.533500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>3</td>\n",
       "      <td>10.456800</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>4</td>\n",
       "      <td>10.555500</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>5</td>\n",
       "      <td>10.552400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>6</td>\n",
       "      <td>10.441400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>7</td>\n",
       "      <td>10.562700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>8</td>\n",
       "      <td>10.534900</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>9</td>\n",
       "      <td>10.361700</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <td>10</td>\n",
       "      <td>10.368100</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table><p>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "109610\n",
      "38573\n",
      "tensor(10.5335, device='cuda:0', grad_fn=<NllLossBackward0>) tensor([[-0.0571, -0.2812,  0.5625,  ...,  0.0903,  0.5742, -0.1523],\n",
      "        [ 0.2246,  0.0074, -0.9375,  ..., -0.1895, -1.1250,  0.8359],\n",
      "        [ 0.3008,  0.8164,  0.0913,  ...,  1.1875,  0.7578, -0.5742],\n",
      "        ...,\n",
      "        [-0.4238, -0.7852, -0.7344,  ..., -0.0815, -0.2246,  0.4238],\n",
      "        [ 0.8633, -0.8125, -0.0226,  ..., -0.2236, -0.5508, -0.6172],\n",
      "        [ 0.1147,  1.0469, -0.3320,  ..., -0.5391,  0.5000, -0.3613]],\n",
      "       device='cuda:0', dtype=torch.bfloat16, grad_fn=<ViewBackward0>) torch.bfloat16 tensor([7081,  313, 8563,  ..., 2228,   15,  436], device='cuda:0') torch.int64\n",
      "273694\n",
      "218013\n",
      "tensor(10.4568, device='cuda:0', grad_fn=<NllLossBackward0>) tensor([[ 1.1484,  0.5625, -0.3027,  ...,  0.6562,  0.3066,  0.1963],\n",
      "        [ 0.0347,  0.0811, -1.6875,  ..., -0.1797,  0.3770,  0.3164],\n",
      "        [ 0.6992, -0.0352,  1.0156,  ..., -0.1196,  0.6602,  0.1357],\n",
      "        ...,\n",
      "        [-0.8594, -0.0137,  0.0115,  ...,  0.7852,  0.8359,  0.5586],\n",
      "        [-0.2490, -0.1904,  0.3418,  ..., -0.0811,  0.2832,  0.1465],\n",
      "        [-0.1475,  1.4922,  0.2217,  ...,  0.6328,  0.2559, -1.0000]],\n",
      "       device='cuda:0', dtype=torch.bfloat16, grad_fn=<ViewBackward0>) torch.bfloat16 tensor([201,  66, 201,  ..., 927,  15, 737], device='cuda:0') torch.int64\n",
      "204334\n",
      "318572\n",
      "tensor(10.5555, device='cuda:0', grad_fn=<NllLossBackward0>) tensor([[ 0.4023,  0.1455, -0.2754,  ..., -0.1011,  0.3691, -0.0771],\n",
      "        [-0.0098, -0.3906,  0.4980,  ...,  0.4434, -0.9727,  0.6953],\n",
      "        [ 0.0308, -0.3047,  0.7188,  ...,  0.0439, -0.5977,  0.1084],\n",
      "        ...,\n",
      "        [ 0.0718, -0.3145, -0.7852,  ..., -0.7734, -0.2305, -1.2812],\n",
      "        [-0.2754,  0.0240, -0.0942,  ..., -1.0859, -0.3965, -0.3887],\n",
      "        [-0.1562,  0.3477,  0.5352,  ..., -0.0869,  0.0098,  0.2178]],\n",
      "       device='cuda:0', dtype=torch.bfloat16, grad_fn=<ViewBackward0>) torch.bfloat16 tensor([ 9924,  8209,  1457,  ..., 15603,    29,   298], device='cuda:0') torch.int64\n",
      "104242\n",
      "261743\n",
      "tensor(10.5524, device='cuda:0', grad_fn=<NllLossBackward0>) tensor([[-0.9727,  0.1680,  0.0981,  ..., -1.1953, -0.1650, -0.1885],\n",
      "        [ 0.2061,  0.3359, -0.1816,  ...,  0.0864, -0.3926, -0.1602],\n",
      "        [ 0.0223, -0.5195,  0.0786,  ...,  1.0078, -0.0728,  0.5312],\n",
      "        ...,\n",
      "        [ 0.5781, -0.4180,  0.6094,  ...,  0.5312, -0.6211,  0.0981],\n",
      "        [-0.0703,  0.1250, -0.2910,  ...,  0.2520, -0.3398,  0.1904],\n",
      "        [-0.0635,  0.5000,  0.6719,  ...,  0.0815,  0.4551, -0.0134]],\n",
      "       device='cuda:0', dtype=torch.bfloat16, grad_fn=<ViewBackward0>) torch.bfloat16 tensor([  344,  4390,  8678,  ...,  1999, 19738,  4748], device='cuda:0') torch.int64\n",
      "235058\n",
      "178647\n",
      "tensor(10.4414, device='cuda:0', grad_fn=<NllLossBackward0>) tensor([[ 0.4297, -0.1279,  0.3203,  ...,  0.3477,  0.8828, -0.1973],\n",
      "        [ 0.0825,  1.2656, -1.2109,  ...,  0.1680, -0.1021,  0.5312],\n",
      "        [-0.2139,  0.5117, -0.9141,  ...,  0.3301, -0.3574, -0.9883],\n",
      "        ...,\n",
      "        [ 0.0679,  0.1641, -0.3203,  ..., -0.7266,  0.8672,  0.4199],\n",
      "        [ 1.2031,  0.2891,  0.1357,  ..., -0.2812,  0.0322,  1.5000],\n",
      "        [-0.7734,  0.5156,  0.0513,  ...,  0.3809,  0.0176, -0.1094]],\n",
      "       device='cuda:0', dtype=torch.bfloat16, grad_fn=<ViewBackward0>) torch.bfloat16 tensor([   77,   201,    66,  ...,  4833,   521, 23351], device='cuda:0') torch.int64\n",
      "225275\n",
      "290107\n",
      "tensor(10.5627, device='cuda:0', grad_fn=<NllLossBackward0>) tensor([[-0.5391,  0.1289, -0.0918,  ..., -0.4160, -0.1196, -0.0542],\n",
      "        [ 0.8633, -0.2031,  0.5664,  ...,  0.2051,  0.7344, -0.0247],\n",
      "        [ 0.7695,  0.4336, -0.2197,  ...,  0.7031,  0.2402, -0.0172],\n",
      "        ...,\n",
      "        [-0.1113, -0.0688, -0.0801,  ..., -0.9961,  0.4883,  0.5039],\n",
      "        [ 0.0457,  0.0581,  0.0835,  ...,  0.1377,  0.0403,  0.5898],\n",
      "        [-0.5508, -0.6445, -0.3594,  ..., -0.4492,  0.6914,  0.9805]],\n",
      "       device='cuda:0', dtype=torch.bfloat16, grad_fn=<ViewBackward0>) torch.bfloat16 tensor([  267, 23724,  1206,  ...,   650, 29570,   628], device='cuda:0') torch.int64\n",
      "105315\n",
      "358768\n",
      "tensor(10.5349, device='cuda:0', grad_fn=<NllLossBackward0>) tensor([[ 0.3945, -0.4590, -0.0364,  ...,  0.4141, -0.0051, -0.1240],\n",
      "        [-0.3320, -0.3008, -0.1455,  ..., -0.0178, -0.0088,  0.2500],\n",
      "        [-0.1562, -0.0649,  0.5039,  ..., -0.2285, -0.5352,  0.5430],\n",
      "        ...,\n",
      "        [ 0.6523,  0.1982,  0.1895,  ...,  0.2520,  0.0297, -0.6328],\n",
      "        [ 0.1396, -0.2617,  0.6641,  ..., -0.6953, -0.6719,  0.2676],\n",
      "        [ 0.2598,  0.7617, -0.9336,  ..., -0.1758, -0.8125, -0.0664]],\n",
      "       device='cuda:0', dtype=torch.bfloat16, grad_fn=<ViewBackward0>) torch.bfloat16 tensor([2582, 2150,  492,  ...,  709,   17, 1999], device='cuda:0') torch.int64\n",
      "332421\n",
      "112040\n",
      "tensor(10.3617, device='cuda:0', grad_fn=<NllLossBackward0>) tensor([[-0.0879, -0.0342,  0.3750,  ..., -0.1768,  0.5391,  1.1016],\n",
      "        [ 0.1309,  0.3340, -1.0078,  ..., -0.7500,  0.2559,  0.2383],\n",
      "        [-0.7812,  1.1172, -0.3555,  ..., -0.1270, -0.1250,  0.5312],\n",
      "        ...,\n",
      "        [ 0.6523,  0.1016, -0.6602,  ..., -0.7969,  0.1455, -0.3848],\n",
      "        [-0.1797, -0.7344, -0.0540,  ..., -0.2129, -0.0312,  0.2773],\n",
      "        [ 0.4609,  1.7891,  0.0962,  ..., -0.1768, -0.7539,  0.0123]],\n",
      "       device='cuda:0', dtype=torch.bfloat16, grad_fn=<ViewBackward0>) torch.bfloat16 tensor([  201,    37, 12861,  ..., 13257,    17,  7103], device='cuda:0') torch.int64\n",
      "198053\n",
      "378413\n",
      "tensor(10.3681, device='cuda:0', grad_fn=<NllLossBackward0>) tensor([[ 0.4688, -1.0312,  0.5352,  ..., -0.4941,  0.9609, -0.7070],\n",
      "        [ 0.5742,  0.7070,  0.2461,  ..., -0.7461,  0.2832,  0.4121],\n",
      "        [-0.9688,  0.2676, -0.4023,  ...,  0.1357,  0.3477,  0.5117],\n",
      "        ...,\n",
      "        [ 1.4375,  0.2002, -0.3164,  ...,  0.4805,  0.4746, -0.3047],\n",
      "        [ 0.1631, -0.4121, -1.0234,  ..., -0.2178,  0.0320,  0.2734],\n",
      "        [ 0.3223, -0.8789,  0.4297,  ...,  1.0469,  0.2207,  0.8945]],\n",
      "       device='cuda:0', dtype=torch.bfloat16, grad_fn=<ViewBackward0>) torch.bfloat16 tensor([ 1122,  8452, 16062,  ...,   201,    51,  7247], device='cuda:0') torch.int64\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Training completed. Do not forget to share your model on huggingface.co/models =)\n",
      "\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TrainOutput(global_step=10, training_loss=10.490837955474854, metrics={'train_runtime': 18.418, 'train_samples_per_second': 1.086, 'train_steps_per_second': 0.543, 'total_flos': 0.0, 'train_loss': 10.490837955474854, 'epoch': 0.0})"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef6f82a9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
